#!/usr/bin/env bash

# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

# Firecracker devtool
#
# Use this script to build and test Firecracker.
#
# TL;DR
# Make sure you have Docker installed and properly configured
# (http://docker.com). Then,
#   building: `./devtool build`
#     Then find the binaries under build/debug/
#   testing: `./devtool test`
#     Will run the entire test battery; will take serveral minutes to complete.
#   deep-dive: `./devtool shell`
#     Open a shell prompt inside the container. Then build or test (or do
#     anything, really) manually.
#
# Still TL;DR: have Docker; ./devtool build; ./devtool test; ./devtool help.
#
#
# Both building and testing are done inside a Docker container. Please make sure
# you have Docker up and running on your system (see http:/docker.com) and your
# user has permission to run Docker containers.
#
# The Firecracker sources dir will be bind-mounted inside the development
# container (under /firecracker) and any files generated by the build process
# will show up under the build/ dir.  This includes the final binaries, as well
# as any intermediate or cache files.
#
# By default, all devtool commands run the container transparently, removing
# it after the command completes. Any persisting files will be stored under
# build/.
# If, for any reason, you want to access the container directly, please use
# `devtool shell`. This will perform the initial setup (bind-mounting the
# sources dir, setting privileges) and will then drop into a BASH shell inside
# the container.
#
# Building:
#   Run `./devtool build`.
#   By default, the debug binaries are built and placed under build/debug/.
#   To build the release version, run `./devtool build --release` instead.
#   You can then find the binaries under build/release/.
#
# Testing:
#   Run `./devtool test`.
#   This will run the entire integration test battery. The testing system is
#   based on pytest (http://pytest.org).
#
# Updating the container image:
#   Run `./devtool update`.
#
# Opening a shell prompt inside the development container:
#   Run `./devtool shell`.
#
# Additional information:
#   Run `./devtool help`.
#
#
# TODO:
#   - Cache test binaries, preserving them across `./devtool test` invocations.
#     At the moment, Firecracker is rebuilt everytime a test is run.
#   - List tests by parsing the `pytest --collect-only` output.
#   - Implement test filtering with `./devtool test --filter <filter>`
#   - Find an easier way to run individual tests on existing binaries.
#   - Add a `./devtool run` command to set up and run Firecracker.
#   - Add a `./devtool diag` command to help with troubleshooting, by checking
#     the most common failure conditions.
#   - Add a `./devtool build-devctr` command to build the development container
#     from its Dockerfile.
#   - Look into caching the Cargo registry within the container and if that
#     would help with reproducible builds (in addition to pinning Cargo.lock)

# Development container image (name:tag)
DEVCTR_IMAGE="fcuvm/dev:latest"

# Naming things is hard
MY_NAME="Firecracker $(basename "$0")"

# Full path to the Firecracker tools dir on the host.
FC_TOOLS_DIR=$(cd "$(dirname $0)" && pwd)

# Full path to the Firecracker sources dir on the host.
FC_ROOT_DIR=$(cd "${FC_TOOLS_DIR}/.." && pwd)

# Full path to the build dir on the host.
FC_BUILD_DIR="${FC_ROOT_DIR}/build"

# Full path to the cargo registry dir on the host. This appears on the host
# because we want to persist the cargo registry across container invocations.
# Otherwise, any rust crates from crates.io would be downloaded again each time
# we build or test.
CARGO_REGISTRY_DIR="${FC_BUILD_DIR}/cargo_registry"

# Full path to the cargo target dir on the host.
CARGO_TARGET_DIR="${FC_BUILD_DIR}/cargo_target"


# Full path to the Firecracker sources dir, as bind-mounted in the container.
CTR_FC_ROOT_DIR="/firecracker"

# Full path to the build dir, as bind-mounted in the container.
CTR_FC_BUILD_DIR="${CTR_FC_ROOT_DIR}/build"

# Full path to the cargo registry dir, as bind-mounted in the container.
CTR_CARGO_REGISTRY_DIR="$CTR_FC_BUILD_DIR/cargo_registry"

# Full path to the cargo target dir, as bind-mounted in the container.
CTR_CARGO_TARGET_DIR="$CTR_FC_BUILD_DIR/cargo_target"

# Full path to the microVM images cache dir
CTR_MICROVM_IMAGES_DIR="$CTR_FC_BUILD_DIR/img"

# Global options received by $0
# These options are not command-specific, so we store them as global vars
OPT_UNATTENDED=false


# Send a decorated message to stdout, followed by a new line
#
say() {
    [ -t 1 ] && [ -n "$TERM" ] \
        && echo "$(tput setaf 2)[$MY_NAME]$(tput sgr0) $@" \
        || echo "[$MY_NAME] $@"
}

# Send a decorated message to stdout, without a trailing new line
#
say_noln() {
    [ -t 1 ] && [ -n "$TERM" ] \
        && echo -n "$(tput setaf 2)[$MY_NAME]$(tput sgr0) $@" \
        || echo "[$MY_NAME] $@"
}

# Send a text message to stderr
#
say_err() {
    [ -t 2 ] && [ -n "$TERM" ] \
        && echo "$(tput setaf 1)[$MY_NAME] $@$(tput sgr0)" 1>&2 \
        || echo "[$MY_NAME] $@" 1>&2
}

# Send a warning-highlighted text to stdout
say_warn() {
    [ -t 1 ] && [ -n "$TERM" ] \
        && echo "$(tput setaf 3)[$MY_NAME] $@$(tput sgr0)" \
        || echo "[$MY_NAME] $@"
}

# Exit with an error message and (optional) code
# Usage: die [-c <error code>] <error message>
#
die() {
    code=1
    [[ "$1" = "-c" ]] && {
        code="$2"
        shift 2
    }
    say_err "$@"
    exit $code
}

# Exit with an error message if the last exit code is not 0
#
ok_or_die() {
    code=$?
    [[ $code -eq 0 ]] || die -c $code "$@"
}

# Check if Docker is available and exit if it's not.
# Upon returning from this call, the caller can be certain Docker is available.
#
ensure_docker() {
    which docker > /dev/null 2>&1
    ok_or_die "Docker not found. Aborting." \
        "Please make sure you have Docker (http://docker.com) installed" \
        "and properly configured."

    docker ps > /dev/null 2>&1
    ok_or_die "Error accessing Docker. Please make sure the Docker daemon" \
        "is running and that you are part of the docker group." \
        "For more information, see" \
        "https://docs.docker.com/install/linux/linux-postinstall/"
}

# Attempt to download our Docker image. Exit if that fails.
# Upon returning from this call, the caller can be certain our Docker image is
# available on this system.
#
ensure_devctr() {

    # We depend on having Docker present.
    ensure_docker

    # Check if we have the container image available locally. Attempt to
    # download it, if we don't.
    [[ $(docker images -q "$DEVCTR_IMAGE" | wc -l) -gt 0 ]] || cmd_update
    ok_or_die "Error pulling docker image. Aborting."
}

# Check if /dev/kvm exists. Exit if it doesn't.
# Upon returning from this call, the caller can be certain /dev/kvm is
# available.
#
ensure_kvm() {
    [[ -c /dev/kvm ]] || die "/dev/kvm not found. Aborting."
}

# Make sure the build/ dirs are available. Exit if we can't create them.
# Upon returning from this call, the caller can be certain the build/ dirs exist.
#
ensure_build_dir() {
    for dir in "$FC_BUILD_DIR" "$CARGO_TARGET_DIR" "$CARGO_REGISTRY_DIR"; do
        mkdir -p "$dir" || die "Error: cannot create dir $dir"
        [ -x "$dir" ] && [ -w "$dir" ] || \
            {
                say "Wrong permissions for $dir. Attempting to fix them ..."
                chmod +x+w "$dir"
            } || \
            die "Error: wrong permissions for $dir. Should be +x+w"
    done
}

# Fix build/ dir permissions after a privileged container run.
# Since the privileged cotainer runs as root, any files it creates will be
# owned by root. This fixes that by recursively changing the ownership of build/
# to the current user.
#
fix_build_dir_perms() {
    # Yes, running Docker to get elevated privileges, just to chown some files
    # is a dirty hack.
    run_devctr \
        -- \
        chown -R "$(id -u):$(id -g)" "$CTR_FC_BUILD_DIR"
}

# Prompt the user for confirmation before proceeding.
# Args:
#   $1  prompt text.
#       Default: Continue? (y/n)
#   $2  confirmation input.
#       Default: y
# Return:
#   exit code 0 for successful confirmation
#   exit code != 0 if the user declined
#
get_user_confirmation() {

    # Pass if running unattended
    [[ "$OPT_UNATTENDED" = true ]] && return 0

    # Fail if STDIN is not a terminal (there's no user to confirm anything)
    [[ -t 0 ]] || return 1

    # Otherwise, ask the user
    #
    msg=$([ -n "$1" ] && echo -n "$1" || echo -n "Continue? (y/n) ")
    yes=$([ -n "$2" ] && echo -n "$2" || echo -n "y")
    say_noln "$msg"
    read c && [ "$c" = "$yes" ] && return 0
    return 1
}

# Helper function to run the dev container.
# Usage: run_devctr <docker args> -- <container args>
# Example: run_devctr --privileged -- bash -c "echo 'hello world'"
run_devctr() {
    docker_args=()
    ctr_args=()
    docker_args_done=false
    while [[ $# -gt 0 ]]; do
        [[ "$1" = "--" ]] && {
            docker_args_done=true
            shift
            continue
        }
        [[ $docker_args_done = true ]] && ctr_args+=("$1") || docker_args+=("$1")
        shift
    done

    # If we're running in a terminal, pass the terminal to Docker and run
    # the container interactively
    [[ -t 0 ]] && docker_args+=("-i")
    [[ -t 1 ]] && docker_args+=("-t")

    # Finally, run the dev container
    #
    docker run "${docker_args[@]}" \
        --init \
        --rm \
        --volume "$FC_ROOT_DIR:$CTR_FC_ROOT_DIR" \
        --env OPT_LOCAL_IMAGES_PATH="$(dirname "$CTR_MICROVM_IMAGES_DIR")" \
        --env PYTHONDONTWRITEBYTECODE=1 \
        "$DEVCTR_IMAGE" "${ctr_args[@]}"
}

# `$0 help`
# Show the detailed devtool usage information.
#
cmd_help() {
    echo ""
    echo "Firecracker $(basename $0)"
    echo "Usage: $(basename $0) [<args>] <command> [<command args>]"
    echo ""
    echo "Global arguments"
    echo "    -y, --unattended         Run unattended. Assume the user would always"
    echo "                             answer \"yes\" to any confirmation prompt."
    echo ""
    echo "Available commands:"
    echo ""
    echo "    build [--debug|--release] [-- [<cargo args>]]"
    echo "        Build the Firecracker binaries."
    echo "        Firecracker is built using the Rust build system (cargo). All arguments after --"
    echo "        will be passed through to cargo."
    echo "        --debug             Build the debug binaries. This is the default."
    echo "        --release           Build the release binaries."
    echo ""
    echo "    help"
    echo "        Display this help message."
    echo ""
    echo "    shell [--privileged]"
    echo "        Launch the development container and open an interactive BASH shell."
    echo "        -p, --privileged    Run the container as root, in privileged mode."
    echo "                            Running Firecracker via the jailer requires elevated"
    echo "                            privileges, though the build phase does not."
    echo ""
    echo "    test [-- [<pytest args>]]"
    echo "        Run the Firecracker integration tests."
    echo "        The Firecracker testing system is based on pytest. All arguments after --"
    echo "        will be passed through to pytest."
    echo ""
    echo "    update"
    echo "        Pull the latest Docker image used to build and test Firecracker."
    echo ""
}

# `$0 build` - build Firecracker
# Please see `$0 help` for more information.
#
cmd_build() {

    # By default, we'll build the debug binaries.
    profile="debug"

    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")  { cmd_help; exit 1;     } ;;
            "--debug")      { profile="debug";      } ;;
            "--release")    { profile="release";    } ;;
            "--")           { break;                } ;;
            *)
                die "Unknown argument: $1. Please use --help for help."
            ;;
        esac
        shift
    done

    # Check prerequisites
    ensure_devctr
    ensure_build_dir

    say "Starting build ($profile) ..."

    # Cargo uses the debug profile by default. If we're building the release
    # binaries, we need to pass an extra argument to cargo.
    cargo_args=("$@")
    [ $profile = "release" ] && cargo_args+=("--release")

    # Run the cargo build process inside the container.
    # We don't need any special privileges for the build phase, so we run the
    # container as the current user/group.
    run_devctr \
        --user "$(id -u):$(id -g)" \
        --workdir "$CTR_FC_ROOT_DIR" \
        -- \
        cargo build \
            --target-dir "${CTR_CARGO_TARGET_DIR}" \
            "${cargo_args[@]}"
    ret=$?

    # If `cargo build` was successful, let's copy the binaries to a more
    # accessible location (under build/release or build/debug).
    if [ $ret -eq 0 ]; then
        mkdir -p "${FC_BUILD_DIR}/${profile}"
        cp -f "${CARGO_TARGET_DIR}/x86_64-unknown-linux-musl/${profile}"/{firecracker,jailer} \
            "${FC_BUILD_DIR}/${profile}/"
        ret=$?
    fi
    [ $ret -eq 0 ] && \
        say "Build complete. Binaries placed under ${FC_BUILD_DIR}/${profile}/"
    return $ret
}

# `$0 test` - run integration tests
# Please see `$0 help` for more information.
#
cmd_test() {

    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")      { cmd_help; exit 1; } ;;
            "--")               { break;            } ;;
            *)
                die "Unknown argument: $1. Please use --help for help."
            ;;
        esac
        shift
    done

    # Check prerequisites.
    ensure_kvm
    ensure_devctr
    ensure_build_dir


    # If we got to here, we've got all we need to continue.
    say "Starting test run ..."

    # Testing (running Firecracker via the jailer) needs root access,
    # in order to set-up the Firecracker jail (manipulating cgroups, net
    # namespaces, etc).
    # We need to run a privileged container to get that kind of access.
    run_devctr \
        --privileged \
        --security-opt seccomp=unconfined \
        --ulimit core=0 \
        --workdir "$CTR_FC_ROOT_DIR/tests" \
        -- \
        pytest "$@"

    ret=$?

    # Running as root would have created some root-owned files under the build
    # dir. Let's fix that.
    fix_build_dir_perms

    return $ret
}

# `$0 shell` - drop to a shell prompt inside the dev container
# Please see `$0 help` for more information.
#
cmd_shell() {

    # By default, we run the container as the current user.
    privileged=false

    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")          { cmd_help; exit 1; } ;;
            "-p"|"--privileged")    { privileged=true;  } ;;
            *)
                die "Unknown argument: $1. Please use --help for help."
            ;;
        esac
        shift
    done

    # Make sure we have what we need to continue.
    ensure_devctr
    ensure_build_dir

    if [[ $privileged = true ]]; then
        # If requested, spin up a privileged container.
        #
        say "Dropping to a privileged shell prompt ..."
        say "Note: $FC_ROOT_DIR is bind-mounted under $CTR_FC_ROOT_DIR"
        say_warn "You are running as root; any files that get created under" \
            "$CTR_FC_ROOT_DIR will be owned by root."
        run_devctr \
            --privileged \
            --security-opt seccomp=unconfined \
            --workdir "$CTR_FC_ROOT_DIR" \
            -- \
            bash
        ret=$?

        # Running as root may have created some root-owned files under the build
        # dir. Let's fix that.
        #
        fix_build_dir_perms
    else
        say "Dropping to shell prompt as user $(whoami) ..."
        say "Note: $FC_ROOT_DIR is bind-mounted under $CTR_FC_ROOT_DIR"
        say_warn "You won't be able to run Firecracker via the jailer," \
            "but you can still build it."
        say "You can use \`$0 shell --privileged\` to get a root shell."

        [ -w /dev/kvm ] || \
            say_warn "WARNING: user $(whoami) doesn't have permission to" \
                "access /dev/kvm. You won't be able to run Firecracker."

        run_devctr \
            --user "$(id -u):$(id -g)" \
            --device=/dev/kvm:/dev/kvm \
            --workdir "$CTR_FC_ROOT_DIR" \
            --env PS1="$(whoami)@\h:\w\$ " \
            -- \
            bash --norc
        ret=$?
    fi

    return $ret
}


# `$0 update` - download the latest dev container image
# Please see `$0 help` for more information.
#
cmd_update() {
    confirmed=false
    while [ $# -gt 0 ]; do
        case "$1" in
            -h|--help)          { cmd_help; exit 1;     } ;;
            -y|--yes)           { confirmed=true;       } ;;
            *)
                die "Unknown argument: $1. Use --help for help."
        esac
        shift
    done

    ensure_docker

    [ $confirmed = true ] || {
        say "About to pull docker image $DEVCTR_IMAGE"
        get_user_confirmation || die "Aborted."
    }

    say "Updating ${DEVCTR_IMAGE} ..."
    docker pull "${DEVCTR_IMAGE}"
}

main() {

    # Parse main command line args.
    #
    while [ $# -gt 0 ]; do
        case "$1" in
            -h|--help)              { cmd_help; exit 1;     } ;;
            -y|--unattended)        { OPT_UNATTENDED=true;  } ;;
            -*)
                die "Unknown arg: $1. Please use \`$0 help\` for help."
            ;;
            *)
                break
            ;;
        esac
        shift
    done

    # $1 is now a command name. Check if it is a valid command and, if so,
    # run it.
    #
    declare -f "cmd_$1" > /dev/null
    ok_or_die "Unknown command: $1. Please use \`$0 help\` for help."

    cmd=cmd_$1
    shift

    # $@ is now a list of command-specific args
    #
    $cmd "$@"
}

main "$@"
